#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
ES事件数据导出到Excel工具
功能：导出指定事件索引中今天的数据到Excel文件
作者：AI助手
日期：2024
"""

import json
import pandas as pd
from datetime import datetime, timedelta
from elasticsearch import Elasticsearch
# from elasticsearch.helpers import scan  # 暂时不使用scan，改用search
import logging
from tqdm import tqdm
import argparse
import sys
import os
import re

# 配置日志
# 确保log目录存在
log_dir = 'log'
if not os.path.exists(log_dir):
    os.makedirs(log_dir)

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler(os.path.join(log_dir, 'es_event_export.log'), encoding='utf-8'),
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger(__name__)


class ESEventExporter:
    def __init__(self, es_host="http://192.168.110.127:9200", index_name="device_event_1928017374610010112_2025-08"):
        """
        初始化ES事件导出器

        Args:
            es_host (str): ES服务地址
            index_name (str): 索引名称
        """
        try:
            self.es = Elasticsearch(
                hosts=[es_host],
                request_timeout=30,  # 减少超时时间
                max_retries=1,       # 减少重试次数
                retry_on_timeout=True,
                # 指定API版本兼容性，解决版本不匹配问题
                headers={'Accept': 'application/vnd.elasticsearch+json; compatible-with=8'}
            )
            logger.info(f"ES客户端初始化成功，服务器: {es_host}")
        except Exception as e:
            logger.error(f"ES客户端初始化失败: {str(e)}")
            raise
        self.index_name = index_name
        self.batch_size = 5000  # 每批处理数量
        
        # 设备信息映射表
        self.device_mapping = {
            '1952695156316139520': {'sn': 'MT22AE50250600012', 'usage': '收货'},
            '1952689993908056064': {'sn': 'MT22AE50250600048', 'usage': '拣货'},
            '1952694000139464704': {'sn': 'MT22AE50250600008', 'usage': '上架'},
            '1928017686301323264': {'sn': 'MT22AE50250600001', 'usage': '定位'},
        }
        
        # 事件标识符到名称的映射表
        self.identifier_mapping = {
            'scan_barcode_1d': '扫描一维条码',
            'scan_barcode_2d': '扫描二维条码',
            'wifi_connect_success': 'wifi连接成功',
            'scan_barcode': '扫描条码',
            'fall_detected': 'PDA跌落',
            'shutdown_reason': '关机原因',
            'scan_barcode_failure': '扫码失败',
            'aoa_label': 'AOA定位',
            'scan_barcode_failure_picture': '扫码失败图片'
        }

    def diagnose_connection(self):
        """诊断ES连接问题"""
        try:
            # 获取集群信息
            info = self.es.info()
            logger.info(f"ES集群信息: {info['name']}, 版本: {info['version']['number']}")
            
            # 获取集群健康状态
            health = self.es.cluster.health()
            logger.info(f"集群状态: {health['status']}, 节点数: {health['number_of_nodes']}")
            
            return True
        except Exception as e:
            logger.error(f"连接诊断失败: {str(e)}")
            return False

    def is_timestamp_field(self, field_name, field_value):
        """
        判断字段是否为时间戳字段
        
        Args:
            field_name (str): 字段名
            field_value: 字段值
            
        Returns:
            bool: 是否为时间戳字段
        """
        # 检查字段名是否包含时间相关关键词
        time_keywords = ['time', 'Time', 'timestamp', 'Timestamp']
        if not any(keyword in field_name for keyword in time_keywords):
            return False
            
        # 检查字段值是否为数字
        try:
            timestamp_value = int(field_value) if isinstance(field_value, str) else field_value
            if not isinstance(timestamp_value, (int, float)):
                return False
        except (ValueError, TypeError):
            return False
            
        # 检查是否为13位数字（Unix毫秒时间戳）
        if len(str(int(timestamp_value))) != 13:
            return False
            
        # 检查时间戳是否在合理范围内（2020年到2030年）
        min_timestamp = 1577836800000  # 2020-01-01 00:00:00 UTC in milliseconds
        max_timestamp = 1893456000000  # 2030-01-01 00:00:00 UTC in milliseconds
        
        return min_timestamp <= timestamp_value <= max_timestamp

    def convert_timestamp_to_readable(self, timestamp_ms):
        """
        将Unix毫秒时间戳转换为可读格式
        
        Args:
            timestamp_ms (int): Unix毫秒时间戳
            
        Returns:
            str: 格式化的时间字符串 "yyyy-MM-dd HH:mm:ss.SSS"
        """
        try:
            # 确保时间戳是数值类型
            if isinstance(timestamp_ms, str):
                timestamp_ms = int(timestamp_ms)
            
            # 转换为datetime对象（UTC时间）
            dt_utc = pd.to_datetime(int(timestamp_ms), unit='ms', utc=True)
            # 转换为北京时间
            dt_beijing = dt_utc.tz_convert('Asia/Shanghai')
            # 格式化为包含毫秒的字符串
            return dt_beijing.strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]  # 保留3位毫秒
        except Exception as e:
            logger.warning(f"时间戳转换失败: {timestamp_ms}, 错误: {str(e)}")
            return str(timestamp_ms)  # 转换失败时返回原值

    def convert_json_timestamps(self, json_obj):
        """
        递归转换JSON对象中的时间戳字段
        
        Args:
            json_obj: JSON对象（dict或list）
            
        Returns:
            转换后的JSON对象
        """
        if isinstance(json_obj, dict):
            converted_obj = {}
            for key, value in json_obj.items():
                if self.is_timestamp_field(key, value):
                    # 转换时间戳字段
                    converted_obj[key] = self.convert_timestamp_to_readable(value)
                elif isinstance(value, (dict, list)):
                    # 递归处理嵌套对象
                    converted_obj[key] = self.convert_json_timestamps(value)
                else:
                    # 保持其他字段不变
                    converted_obj[key] = value
            return converted_obj
        elif isinstance(json_obj, list):
            return [self.convert_json_timestamps(item) for item in json_obj]
        else:
            return json_obj

    def convert_content_timestamps(self, content_str):
        """
        转换content字段中的时间戳
        
        Args:
            content_str (str): content字段的JSON字符串
            
        Returns:
            str: 转换后的JSON字符串
        """
        if not content_str or pd.isna(content_str) or content_str == '':
            return content_str
            
        try:
            # 解析JSON字符串
            json_obj = json.loads(content_str)
            
            # 转换其中的时间戳字段
            converted_obj = self.convert_json_timestamps(json_obj)
            
            # 重新序列化为JSON字符串
            return json.dumps(converted_obj, ensure_ascii=False, separators=(',', ':'))
            
        except json.JSONDecodeError:
            # 如果不是有效的JSON，返回原值
            logger.warning(f"content字段不是有效的JSON格式: {content_str[:100]}...")
            return content_str
        except Exception as e:
            logger.warning(f"content时间戳转换失败: {str(e)}")
            return content_str

    def get_today_range(self):
        """
        获取今天的时间范围（Unix毫秒时间戳）

        Returns:
            tuple: (today_start_ms, now_ms)
        """
        # 获取本地时间
        local_now = datetime.now()
        # 今天开始时间（00:00:00）
        today_start = local_now.replace(hour=0, minute=0, second=0, microsecond=0)
        
        # 转换为Unix毫秒时间戳
        today_start_ms = int(today_start.timestamp() * 1000)
        now_ms = int(local_now.timestamp() * 1000)

        logger.info(f"导出时间范围: {today_start.strftime('%Y-%m-%d %H:%M:%S')} ({today_start_ms}) 至 {local_now.strftime('%Y-%m-%d %H:%M:%S')} ({now_ms})")
        return today_start_ms, now_ms

    def build_query(self):
        """
        构建查询条件

        Returns:
            dict: ES查询DSL
        """
        today_start, now = self.get_today_range()

        query = {
            "query": {
                "bool": {
                    "must": [
                        {
                            "range": {
                                "timestamp": {  # 根据您的映射，使用timestamp字段
                                    "gte": today_start,
                                    "lte": now
                                }
                            }
                        }
                    ]
                }
            },
            "_source": {  # 只返回需要的字段，提高性能
                "includes": [
                    "deviceId", "identifier", "content", "createTime"
                ]
            }
        }
        return query

    def fetch_data(self, max_docs=None):
        """
        从ES获取数据

        Args:
            max_docs (int): 最大导出文档数（用于测试）

        Returns:
            list: 文档列表
        """
        try:
            # 测试连接
            logger.info(f"尝试连接ES服务器...")
            ping_result = self.es.ping()
            if not ping_result:
                # 尝试诊断连接问题
                logger.error("ping失败，尝试详细诊断...")
                if not self.diagnose_connection():
                    raise ConnectionError("无法连接到Elasticsearch服务器")
            
            logger.info("ES连接成功")
            # 执行连接诊断
            self.diagnose_connection()

            # 验证索引是否存在
            logger.info(f"检查索引是否存在: {self.index_name}")
            if not self.es.indices.exists(index=self.index_name):
                raise ValueError(f"索引不存在: {self.index_name}")
            
            logger.info("索引验证成功")

            # 构建查询
            query = self.build_query()

            # 估算总数（用于进度条）
            try:
                count_resp = self.es.count(index=self.index_name, query=query['query'])
                total_docs = count_resp['count']
            except Exception as e:
                logger.warning(f"获取文档总数失败，使用默认值: {str(e)}")
                total_docs = 10000  # 默认值

            if max_docs:
                total_docs = min(total_docs, max_docs)

            logger.info(f"预计导出 {total_docs} 条记录")

            # 使用search_after方式获取数据（突破10000条限制）
            docs = []
            search_after = None
            
            # 使用tqdm显示进度
            with tqdm(total=total_docs, desc="导出进度", unit="条") as pbar:
                while not max_docs or len(docs) < max_docs:
                    # 计算本次查询的大小
                    remaining = total_docs - len(docs) if total_docs > len(docs) else self.batch_size
                    current_size = min(self.batch_size, remaining)
                    if max_docs:
                        current_size = min(current_size, max_docs - len(docs))
                    
                    if current_size <= 0:
                        break
                    
                    # 构建搜索体
                    search_body = {
                        "query": query['query'],
                        "_source": query['_source'],
                        "size": current_size,
                        "sort": [
                            {"timestamp": {"order": "asc"}}  # 只使用timestamp排序
                        ]
                    }
                    
                    # 如果有search_after值，添加到查询中
                    if search_after:
                        search_body["search_after"] = search_after
                    
                    try:
                        result = self.es.search(index=self.index_name, body=search_body)
                        hits = result['hits']['hits']
                        
                        if not hits:
                            logger.info("没有更多数据")
                            break
                            
                        # 提取数据
                        for hit in hits:
                            docs.append(hit['_source'])
                            pbar.update(1)
                        
                        # 更新search_after为最后一条记录的排序值
                        if hits:
                            last_hit = hits[-1]
                            search_after = last_hit['sort']
                        
                        # 记录进度
                        if len(docs) % 1000 == 0:
                            logger.info(f"已处理 {len(docs)} 条记录")
                            
                    except Exception as e:
                        logger.error(f"搜索数据时出错: {str(e)}")
                        # 如果是result window错误，提示用户
                        if "result window is too large" in str(e).lower():
                            logger.error("遇到ES结果窗口限制，已自动切换到search_after方式")
                        break

            logger.info(f"数据获取完成，共 {len(docs)} 条记录")
            return docs

        except Exception as e:
            logger.error(f"数据获取失败: {str(e)}")
            raise

    def process_data(self, docs):
        """
        处理数据（类型转换、格式化等）

        Args:
            docs (list): 原始文档列表

        Returns:
            pandas.DataFrame: 处理后的数据框
        """
        if not docs:
            logger.warning("没有数据需要处理")
            return pd.DataFrame()

        # 转换为DataFrame
        df = pd.DataFrame(docs)

        # 数据类型优化 - 只处理需要的字段
        type_mapping = {
            'deviceId': 'string',
            'identifier': 'string',
            'content': 'string'
        }
        
        # 处理缺失值
        df = df.fillna('')

        # 添加固定的type列
        df['type'] = 'other'
        logger.info("已添加type列，固定值为'other'")

        # 处理时间字段（从Unix毫秒时间戳转换，考虑时区）
        time_fields = ['createTime']
        for field in time_fields:
            if field in df.columns:
                try:
                    # 转换Unix毫秒时间戳为UTC时间
                    df[field] = pd.to_datetime(df[field], unit='ms', utc=True, errors='coerce')
                    # 转换为北京时间（UTC+8）
                    df[field] = df[field].dt.tz_convert('Asia/Shanghai')
                    # 格式化为 yyyy-MM-dd HH:mm:ss 格式（去掉时区信息）
                    df[field] = df[field].dt.strftime('%Y-%m-%d %H:%M:%S')
                    logger.info(f"时间字段 {field} 已转换为北京时间格式")
                except Exception as e:
                    logger.warning(f"时间字段 {field} 转换失败: {str(e)}")
                    # 如果时区转换失败，尝试简单转换
                    try:
                        df[field] = pd.to_datetime(df[field], unit='ms', errors='coerce')
                        df[field] = df[field].dt.strftime('%Y-%m-%d %H:%M:%S')
                        logger.warning(f"时间字段 {field} 使用UTC时间（未转换时区）")
                    except Exception as e2:
                        logger.error(f"时间字段 {field} 完全转换失败: {str(e2)}")

        # 应用类型映射
        for col, dtype in type_mapping.items():
            if col in df.columns:
                try:
                    if dtype == 'string':
                        df[col] = df[col].astype('string')
                except Exception as e:
                    logger.warning(f"字段 {col} 类型转换失败: {str(e)}")
        
        # 添加设备信息映射（sn和usage列）
        if 'deviceId' in df.columns:
            logger.info("开始添加设备信息映射...")
            
            # 添加sn列
            df['sn'] = df['deviceId'].map(lambda x: self.device_mapping.get(str(x), {}).get('sn', ''))
            
            # 添加usage列  
            df['usage'] = df['deviceId'].map(lambda x: self.device_mapping.get(str(x), {}).get('usage', ''))
            
            # 统计映射结果
            mapped_devices = df[df['sn'] != '']['deviceId'].nunique()
            total_devices = df['deviceId'].nunique()
            logger.info(f"设备信息映射完成: {mapped_devices}/{total_devices} 个设备有映射信息")
        
        # 添加事件标识符名称映射（name列）
        if 'identifier' in df.columns:
            logger.info("开始添加事件标识符名称映射...")
            
            # 添加name列
            df['name'] = df['identifier'].map(lambda x: self.identifier_mapping.get(str(x), ''))
            
            # 统计映射结果
            mapped_identifiers = df[df['name'] != '']['identifier'].nunique()
            total_identifiers = df['identifier'].nunique()
            logger.info(f"事件标识符映射完成: {mapped_identifiers}/{total_identifiers} 个标识符有映射信息")
        
        # 转换content字段中的时间戳
        if 'content' in df.columns:
            logger.info("开始转换content字段中的时间戳...")
            
            # 统计转换前的情况
            non_empty_content = df[df['content'].notna() & (df['content'] != '')].shape[0]
            logger.info(f"发现 {non_empty_content} 条记录包含content数据")
            
            # 应用时间戳转换
            df['content'] = df['content'].apply(self.convert_content_timestamps)
            
            logger.info("content字段时间戳转换完成")
        
        # 按指定顺序排列列：deviceId, sn, usage, identifier, name, content, type, createTime
        desired_order = ['deviceId', 'sn', 'usage', 'identifier', 'name', 'content', 'type', 'createTime']
        existing_cols = [col for col in desired_order if col in df.columns]
        df = df[existing_cols]

        logger.info(f"数据处理完成，共 {len(df)} 行，{len(df.columns)} 列")
        return df

    def export_to_excel(self, df, output_file=None):
        """
        导出数据到Excel

        Args:
            df (pandas.DataFrame): 数据框
            output_file (str): 输出文件名

        Returns:
            str: 输出文件路径
        """
        if df.empty:
            logger.warning("数据为空，未生成Excel文件")
            return None

        # 生成默认文件名
        if not output_file:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            output_file = f"es_event_export_{self.index_name}_{timestamp}.xlsx"

        try:
            # 写入Excel
            with pd.ExcelWriter(output_file, engine='openpyxl') as writer:
                # 主数据表
                df.to_excel(writer, sheet_name='设备事件数据', index=False)

                # 统计信息表
                local_now = datetime.now()
                today_start = local_now.replace(hour=0, minute=0, second=0, microsecond=0)
                # 计算有映射信息的设备数量
                mapped_device_count = len(df[df['sn'] != '']['deviceId'].unique()) if 'sn' in df.columns else 0
                
                # 计算有映射信息的标识符数量
                mapped_identifier_count = len(df[df['name'] != '']['identifier'].unique()) if 'name' in df.columns else 0
                
                stats_data = {
                    '统计项': [
                        '总记录数',
                        '导出时间',
                        '源索引',
                        '数据时间段',
                        '设备数量',
                        '有映射信息的设备数',
                        '事件标识符数量',
                        '有映射信息的标识符数',
                        '输出字段'
                    ],
                    '值': [
                        len(df),
                        local_now.strftime("%Y-%m-%d %H:%M:%S"),
                        self.index_name,
                        f"{today_start.strftime('%Y-%m-%d %H:%M:%S')} 至 {local_now.strftime('%Y-%m-%d %H:%M:%S')}",
                        df['deviceId'].nunique() if 'deviceId' in df.columns else 0,
                        mapped_device_count,
                        df['identifier'].nunique() if 'identifier' in df.columns else 0,
                        mapped_identifier_count,
                        ', '.join(df.columns.tolist())
                    ]
                }
                stats_df = pd.DataFrame(stats_data)
                stats_df.to_excel(writer, sheet_name='导出统计', index=False)

            logger.info(f"Excel文件已生成: {output_file}")
            logger.info(f"文件大小: {os.path.getsize(output_file) / 1024 / 1024:.2f} MB")
            return output_file

        except Exception as e:
            logger.error(f"Excel导出失败: {str(e)}")
            raise

    def export(self, output_file=None, max_docs=None):
        """
        完整导出流程

        Args:
            output_file (str): 输出文件名
            max_docs (int): 最大导出文档数

        Returns:
            str: 输出文件路径
        """
        try:
            logger.info(f"开始导出事件索引 [{self.index_name}] 今日数据...")

            # 1. 获取数据
            docs = self.fetch_data(max_docs)

            # 2. 处理数据
            df = self.process_data(docs)

            # 3. 导出Excel
            output_path = self.export_to_excel(df, output_file)

            logger.info("事件数据导出完成！")
            return output_path

        except Exception as e:
            logger.error(f"导出过程出错: {str(e)}")
            raise


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='导出ES事件索引今日数据到Excel')
    parser.add_argument('--host', default='http://192.168.110.127:9200', help='ES服务地址')
    parser.add_argument('--index', default='device_event_1928017374610010112_2025-08', help='事件索引名称')
    parser.add_argument('--output', help='输出文件名')
    parser.add_argument('--max-docs', type=int, help='最大导出文档数（测试用）')
    parser.add_argument('--test', action='store_true', help='测试模式（仅导出100条）')

    args = parser.parse_args()

    # 测试模式
    if args.test:
        args.max_docs = 100
        logger.info("测试模式：仅导出前100条记录")

    try:
        # 创建导出器
        exporter = ESEventExporter(es_host=args.host, index_name=args.index)

        # 执行导出
        output_file = exporter.export(
            output_file=args.output,
            max_docs=args.max_docs
        )

        if output_file:
            print(f"\n✅ 事件数据导出成功！")
            print(f"📁 文件路径: {os.path.abspath(output_file)}")
            print(f"📊 记录数量: {pd.read_excel(output_file).shape[0]} 条")
        else:
            print("⚠️ 没有事件数据需要导出")

    except Exception as e:
        logger.error(f"导出失败: {str(e)}")
        sys.exit(1)


if __name__ == "__main__":
    main()